from selenium import webdriver
import ddddocr
from selenium.webdriver.common.by import By
from selenium.webdriver.support.wait import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC

from tools.get_driver import GetDriver


class GetVerifyCode:
    driver = None

    def get_code(self,url):
        try:
            # 1. 通过GetDriver获取浏览器驱动
            self.driver = GetDriver.get_driver(url)

            # 2.等待验证码图片加载
            captcha_img = WebDriverWait(self.driver, 10).until(
                EC.presence_of_element_located((By.CSS_SELECTOR, "img[src*='userverifyentry']"))
            )

            # 3. 对验证码图片进行截图
            captcha_screenshot = captcha_img.screenshot_as_png

            # 4. 初始化ddddocr识别器，识别验证码图片数据
            ocr = ddddocr.DdddOcr()
            verify_code = ocr.classification(captcha_screenshot)

            # 5. 提取浏览器当前会话的Cookie
            cookies = {cookie['name']: cookie['value'] for cookie in self.driver.get_cookies()}

            # 6. 返回识别成功的验证码和Cookie
            return verify_code, cookies

        except Exception as e:
            print(f"处理失败：{str(e)}")
            return None, None
        finally:
            if self.driver:
                GetDriver.quit_driver()
